import numpy as np
from sfepy.discrete.fem import Mesh, FEDomain, Field
from sfepy.discrete import Problem, Functions, Equations, Region
from sfepy.base.base import assert_, output, Struct
from sfepy.solvers import Solver
from sfepy.mechanics.matcoefs import stiffness_from_lame
from sfepy.discrete import (FieldVariable, Material, Integral, Integrals,
                            Equation, Equations, Problem)
from sfepy.mechanics.matcoefs import stiffness_from_youngpoisson
from sfepy.terms import Term
from sfepy.discrete.conditions import Conditions, EssentialBC
import scipy.sparse.linalg as sla
import matplotlib.pyplot as plt
import pandas as pd
import os
import crackgenerate
import itertools
from getfiles import get_files_by_extension, shape_plot


N_EIGS = 5  #总共分析前5阶模态振型




def modal_analysis(eq1: Equation, eq2: Equation, pb: Problem, cavity, n_eigs, eig_solver):
    # 7. 定义方程
    mtx_k = eq1.evaluate_withcavity(mode='weak', dw_mode='matrix', asm_obj=pb.mtx_a, cavity=cavity)
    mtx_m = mtx_k.copy()
    mtx_m.data[:] = 0.0
    mtx_m = eq2.evaluate_withcavity(mode='weak', dw_mode='matrix', asm_obj=mtx_m, cavity=cavity)
    n_rbm = 0
    try:
        eigs, svecs = eig_solver(mtx_k, mtx_m, n_eigs + n_rbm,
                                    eigenvectors=True)
    except sla.ArpackNoConvergence as ee:
        eigs = ee.eigenvalues
        svecs = ee.eigenvectors
        output('only %d eigenvalues converged!' % len(eigs))

    eigs = eigs[n_rbm:]
    svecs = svecs[:, n_rbm:]
    omegas = np.sqrt(eigs)
    freqs = omegas / (2 * np.pi)
    output('number |         eigenvalue |  angular frequency '
            '|          frequency')
    for ii, eig in enumerate(eigs):
        output('%6d | %17.12e | %17.12e | %17.12e'
                % (ii + 1, eig, omegas[ii], freqs[ii]))

    # 计算振型向量
    variables = pb.set_default_state()

    vecs = np.empty((variables.di.n_dof_total, svecs.shape[1]),
                    dtype=np.float64)
    for ii in range(svecs.shape[1]):
        vecs[:, ii] = variables.make_full_vec(svecs[:, ii])

    return freqs, vecs

def cellinfo_generate(domain: FEDomain, omega: Region, 
                      vecs: np.ndarray, cavity: np.ndarray):
    """
    将结构模态振动分析的结果保存为ndarray的形式。
    
    参数
    ---------
    domain : FEDomain
        分析的几何结构的FEDomain对象。
    vecs : ndarray
        模态分析得到的振型向量。
    omega : Region
        表征domain全局区域的Region，通常通过domain.create_region('Omega', 'all')获取。
    cavity : 
        表示模态分析中刚度和质量被置0的单元索引。
    
    返回值
    ---------
    cellinfo : ndarray
        cellinfo.shape = (cellnum, 8, 19)
        cellinfo的第0维对应domain中的不同单元。
        cellinfo的第1为对应六面体单元中的8个节点。
        cellinfo第2维的0位对应每个单元的节点id。
        cellinfo第2维的1~3位对应每个单元的节点(x,y,z)坐标。
        cellinfo第2维的4~6位对应每个单元的节点在第1阶模态下对应的(x,y,z)位移。
                    7~9位对应每个单元的节点在第2阶模态下对应的(x,y,z)位移。
                    ...
                    16~18位对应每个单元的节点在第5阶模态下对应的(x,y,z)位移。
    
    cellstate : ndarray
        cellstate.shape = (cellnum, )
        每一位表示domain中的每个单元的质量和刚度是否被置0。
    """
    cells_nodes = domain.get_conn()
    node_corr = omega.cmesh.coors
    cells_nodes_corr = node_corr[cells_nodes]
    node_vecs = vecs.reshape(-1, 3 * vecs.shape[1])
    cells_nodes_vec = node_vecs[cells_nodes, :]
    cellinfo = np.zeros([cells_nodes.shape[0], 8, 19], dtype=np.float32)
    cellstate = np.ones(cells_nodes.shape[0], dtype=bool)
    cellinfo[:, :, 0] = cells_nodes
    cellinfo[:, :, 1:4] = cells_nodes_corr
    cellinfo[:, :, 4:19] = cells_nodes_vec
    cellstate[cavity] = False

    return cellinfo, cellstate


# 1. 加载网格
modelgeo_path = os.path.join(os.getcwd(), "datagen", "modelgeo")
filesnames = get_files_by_extension(modelgeo_path, '.msh')

output_dir = os.path.join(os.getcwd(), 'datagen', 'dataset')

dim = 3
axis = 1
axis_fix = 0  # 沿x轴方向固定

for filename in filesnames:
    mesh = Mesh.from_file(os.path.join(modelgeo_path, filename))
    domain = FEDomain('domain', mesh)
    bbox = domain.get_mesh_bounding_box()
    min_coor, max_coor = bbox[:, axis]
    min_side_coor, max_side_coor = bbox[:, axis_fix]
    eps = 1e-8 * (max_coor - min_coor)
    eps_side = 1e-8 * (max_side_coor - min_side_coor)
    ax = 'xyz'[:dim][axis]
    aside = 'xyz'[:dim][axis_fix]

    omega = domain.create_region('Omega', 'all')
    left_side = domain.create_region('Rightside',
                                    'vertices in (%s < %.10f)'
                                    % (aside, min_side_coor + eps_side),
                                    'facet')
    right_side = domain.create_region('Rightside',
                                    'vertices in (%s > %.10f)'
                                    % (aside, max_side_coor - eps_side),
                                    'facet')

    # 生成裂纹位置
    # cavity_list = []
    crack_y_list = np.arange(np.abs(4/6 * (max_coor - min_coor)) // 6) * 6 + (5/6 * min_coor + 1/6 * max_coor)
    crack_x_list = np.arange(np.abs(1/2 * (max_side_coor - min_side_coor)) // 6) * 6 + (3/4 * min_side_coor + 1/4 * max_side_coor)
    crack_dir_list = np.array([-1, 1])
    combinations = list(itertools.product(crack_y_list, crack_x_list, crack_dir_list))
    target_corr_z = 0.1
    sim_idx = 0
    for target_corr_y, target_corr_x, crack_dir in combinations:
        target_corr = np.array([target_corr_x, target_corr_y, target_corr_z])
        if crack_dir > 0:
            cavity = crackgenerate.get_cavity_region(domain, omega, 
                                                     target_corr, right_side, crack_dir)
        else:
            cavity = crackgenerate.get_cavity_region(domain, omega, 
                                                     target_corr, left_side, crack_dir)
        # cavity_list.append(cavity)

        mesh = Mesh.from_file(os.path.join(modelgeo_path, filename))
        domain = FEDomain('domain', mesh)
        bbox = domain.get_mesh_bounding_box()
        min_coor, max_coor = bbox[:, axis]
        min_side_coor, max_side_coor = bbox[:, axis_fix]
        eps = 1e-8 * (max_coor - min_coor)
        eps_side = 1e-8 * (max_side_coor - min_side_coor)
        ax = 'xyz'[:dim][axis]
        aside = 'xyz'[:dim][axis_fix]
        omega = domain.create_region('Omega', 'all')
        # 设置约束区域
        fix_bottom = domain.create_region('Bottom',
                                    'vertices in (%s < %.10f)'
                                    % (ax, min_coor + eps),
                                    'facet')

        # 2. 定义材料属性（钢为例）
        youngs_modulus = 210e3  # MPa (mm 量纲)
        poissons_ratio = 0.3
        density = 7.85e-6       # kg/mm³

        # 3. 定义求解器
        eig_conf = Struct(name='evp', kind='eig.scipy', method='eig')

        # 4. 定义有限元场（六面体线性单元）
        order = 1
        field = Field.from_args('fu', np.float64, 'vector', omega,
                                    approx_order=order)
        u = FieldVariable('u', 'unknown', field)
        v = FieldVariable('v', 'test', field, primary_var_name='u')

        # 5. 定义材料
        mtx_d = stiffness_from_youngpoisson(dim, young=youngs_modulus, poisson=poissons_ratio)
        m = Material('m', D=mtx_d, rho=density)
        integral = Integral('i', order=2*order)
        t1 = Term.new('dw_lin_elastic(m.D, v, u)', integral, omega, m=m, v=v, u=u)
        t2 = Term.new('dw_dot(m.rho, v, u)', integral, omega, m=m, v=v, u=u)
        eq1 = Equation('stiffness', t1)
        eq2 = Equation('mass', t2)
        lhs_eqs = Equations([eq1, eq2])
        pb = Problem('modal', equations=lhs_eqs)

        # 6. 定义边界条件
        fixed = EssentialBC('Fixed', fix_bottom, {'u.all' : 0.0})
        pb.time_update(ebcs=Conditions([fixed]))
        pb.update_materials()
        eig_solver = Solver.any_from_conf(eig_conf)
        freqs, vecs = modal_analysis(eq1, eq2, pb, cavity, N_EIGS, eig_solver)
        cellinfo, cellstate = cellinfo_generate(domain, omega, vecs, cavity)
        shape_plot(domain, vecs, 5, freqs)
        np.savez(os.path.join(output_dir, filename[:-4] + 'sim' + str(sim_idx) + '.npz'), 
                 cellinfo=cellinfo, cellstate=cellstate, freqs=freqs)
        sim_idx += 1







